from pyspark.sql.session import SparkSession
from pyspark.sql.functions import *
from pyspark.ml.linalg import Vectors

from pyspark.ml.classification import LogisticRegression

# 创建环境
spark = SparkSession.builder.getOrCreate()

# 1 读取数据
data = spark.read.format("libsvm").load("../../data/人体指标.txt")

data.printSchema()
data.show(truncate=False)

# 2 将数据拆分成训练集和测试集
# 训练集用于训练模型
# 测试集用于测试模型的准确率

train, test = data.randomSplit([0.7, 0.3])

# 3 选择算法
# LogisticRegression：逻辑回归，用于分类算法
lr = LogisticRegression()

# 4 将训练集带入算法训练模型
model = lr.fit(train)

# 5 将测试机带入模型测试模型的准确率
test_predict = model.transform(test)

test_predict.show(truncate=False)

# 准确率 Accuracy
Accuracy = test_predict.where(col("prediction") == col("label")).count() / test_predict.count()
print(f"准确率：{Accuracy}")

# 对于二分类问题需要根据精确度和召回率判断

# 精确度 Precision
# TP/TP+FP
Precision = test_predict.where((col("label") == 1) & (col("prediction") == 1)).count() / test_predict.where(
    col("prediction") == 1).count()

print(f"精确度：{Precision}")

# 召回率 Recall
# TP/TP+FN
Recall = test_predict.where((col("label") == 1) & (col("prediction") == 1)).count() / test_predict.where(
    col("label") == 1).count()
print(f"召回率：{Recall}")

# F1,只有当准确率和召回率都高时才会高
F1 = 2 * Precision * Recall / (Precision + Recall)
print(f"F1：{F1}")
